# utils.py
"""
存放项目通用的辅助函数，用于将 RagDataset 转换为不同的微调格式。
"""
import json
import logging
from typing import List, Dict, Any
import pandas as pd
# from llama_index.core.llama_dataset import RagDataset # 如果RagDataset是自定义类型，调用时需确保可导入

# =====================================================================================
# 函数 1: 将 RagDataset 转换为 SFTTrainer 的 'messages' 格式
# =====================================================================================
def format_rag_for_sfttrainer(rag_dataset) -> List[Dict[str, Any]]:
    """
    将 RagDataset 转换为 SFTTrainer 偏好的 'messages' 对话格式。
    输出格式: [{"messages": [{"role": "user", ...}, {"role": "assistant", ...}]}]

    Args:
        rag_dataset (RagDataset): LlamaIndex 的 RagDataset 对象。

    Returns:
        List[Dict[str, Any]]: 格式化后的数据集列表。
    """
    logging.info("正在将 RagDataset 转换为 SFT 'messages' 对话格式...")
    df = rag_dataset.to_pandas()

    def create_messages_format(row: pd.Series) -> Dict[str, Any]:
        """为DataFrame的一行创建 'messages' 字典"""
        return {
            "messages": [
                {
                    "role": "user",
                    "content": row['query']
                },
                {
                    "role": "assistant",
                    "content": row['reference_answer']
                }
            ]
        }

    # 应用格式化函数
    sft_data = df[['query', 'reference_answer']].apply(create_messages_format, axis=1).tolist()
    logging.info(f"格式转换完成 ('messages' 格式)，共生成 {len(sft_data)} 条对话数据。")
    return sft_data


# =====================================================================================
# 函数 2: 将 RagDataset 转换为 'conversation' 格式
# =====================================================================================
def format_rag_to_xtuner(rag_dataset) -> List[Dict[str, Any]]:
    """
    将 RagDataset 转换为 'conversation' 对话格式。
    输出格式: [{"conversation": [{"input": ..., "output": ...}]}]

    Args:
        rag_dataset (RagDataset): LlamaIndex 的 RagDataset 对象。

    Returns:
        List[Dict[str, Any]]: 格式化后的数据集列表。
    """
    logging.info("正在将 RagDataset 转换为 'conversation' 对话格式...")
    df = rag_dataset.to_pandas()

    def create_conversation_format(row: pd.Series) -> Dict[str, Any]:
        """为DataFrame的一行创建 'conversation' 字典"""
        return {
            "conversation": [
                {
                    "input": row['query'],
                    "output": row['reference_answer']
                }
            ]
        }

    # 应用格式化函数
    sft_data = df[['query', 'reference_answer']].apply(create_conversation_format, axis=1).tolist()
    logging.info(f"格式转换完成 ('conversation' 格式)，共生成 {len(sft_data)} 条对话数据。")
    return sft_data


# =====================================================================================
# 函数 3: 将数据保存为 JSON 文件 (保持不变)
# =====================================================================================
def save_json(data: List[Dict[str, Any]], file_path: str):
    """
    将数据以 JSON 格式保存到文件。

    Args:
        data (List[Dict[str, Any]]): 要保存的数据。
        file_path (str): 输出文件路径。
    """
    logging.info(f"正在将数据集保存到 {file_path}...")
    try:
        with open(file_path, 'w', encoding='utf-8') as f:
            json.dump(data, f, indent=4, ensure_ascii=False)
        logging.info(f"数据集已成功保存。")
    except IOError as e:
        logging.error(f"保存文件时出错: {e}")
        raise